# ======================================================================
# Copyright CERFACS (March 2019)
# Contributor: Adrien Suau (adrien.suau@cerfacs.fr)
#
# This software is governed by the CeCILL-B license under French law and
# abiding  by the  rules of  distribution of free software. You can use,
# modify  and/or  redistribute  the  software  under  the  terms  of the
# CeCILL-B license as circulated by CEA, CNRS and INRIA at the following
# URL "http://www.cecill.info".
#
# As a counterpart to the access to  the source code and rights to copy,
# modify and  redistribute granted  by the  license, users  are provided
# only with a limited warranty and  the software's author, the holder of
# the economic rights,  and the  successive licensors  have only limited
# liability.
#
# In this respect, the user's attention is drawn to the risks associated
# with loading,  using, modifying and/or  developing or reproducing  the
# software by the user in light of its specific status of free software,
# that  may mean  that it  is complicated  to manipulate,  and that also
# therefore  means that  it is reserved for  developers and  experienced
# professionals having in-depth  computer knowledge. Users are therefore
# encouraged  to load and  test  the software's  suitability as  regards
# their  requirements  in  conditions  enabling  the  security  of their
# systems  and/or  data to be  ensured and,  more generally,  to use and
# operate it in the same conditions as regards security.
#
# The fact that you  are presently reading this  means that you have had
# knowledge of the CeCILL-B license and that you accept its terms.
# ======================================================================

r"""Arithmetic quantum routines implementation.

This module implements a function to add a constant to a quantum register and
a function to subtract a constant from a quantum register.

.. warning::
   At the moment this module is just a wrapper around qat. This will probably
   change in the future as the oracles implementation needs to have additional
   pre- and post-conditions that are not listed in the documentation (behaviour
   with respect to overflow / underflow, precision, gate number, ...).
"""

from qat.lang.AQASM import qftarith, classarith
from qat.lang.AQASM.gates import X, CNOT, CCNOT
from qat.lang.AQASM.routines import QRoutine
from qat.lang.AQASM.misc import build_gate


@build_gate("add_const", [int, int], arity=lambda n, num: n)
def add_const_qft(n: int, number: int) -> QRoutine:
    r"""Add a constant to the given :math:`n`\ -qubit register.

    The following properties are desirable ([ ] means that the property is not
    guaranteed to be fulfilled, [x] guarantees that the property is fulfilled
    by the current implementation):

    - [x] Looping overflow.
    - [ ] Exact results (no non-valid state with very small amplitudes
      appearing).
    - [x] Gate number increase (or decrease) with the number of '1' in the binary
      representation of the given constant.

    :param n: The size of the quantum register we will add `number` to.
    :param number: The number to add to the given quantum register.
    :return: a routine adding number to the given quantum register.
    """
    add_routine = QRoutine()
    qubits = add_routine.new_wires(n)
    # Swap qubits endianness because qat's adder is in LSB
    add_routine.apply(qftarith.add_const(n, number), qubits[::-1])
    return add_routine


@build_gate("sub_const", [int, int], arity=lambda n, num: n)
def sub_const_qft(n: int, number: int) -> QRoutine:
    r"""Subtract a constant to the given `n`\ -qubit register.

    The following properties are desirable ([ ] means that the property is not
    guaranteed to be fulfilled, [x] guarantees that the property is fulfilled
    by the current implementation):

    - [x] Looping underflow.
    - [ ] Exact results (no non-valid state with very small amplitudes
      appearing).
    - [x] Gate number increase (or decrease) with the number of '1' in the binary
      representation of the given constant.

    :param n: The size of the quantum register we will subtract `number` from.
    :param number: The number to subtract from the given quantum register.
    :return: a routine subtracting number from the given quantum register.
    """
    routine = QRoutine()
    sub_const_gate = add_const_qft(n, number).dag()
    qubits = routine.new_wires(sub_const_gate.arity)
    routine.apply(sub_const_gate, qubits)
    return routine


@build_gate("add_const", [int, int], arity=lambda n, num: n)
def add_const_carry(n: int, number: int) -> QRoutine:
    r"""Add a constant to the given :math:`n`\ -qubit register.

    The following properties are desirable ([ ] means that the property is not
    guaranteed to be fulfilled, [x] guarantees that the property is fulfilled
    by the current implementation):

    - [x] Looping overflow.
    - [ ] Exact results (no non-valid state with very small amplitudes
      appearing).
    - [x] Gate number increase (or decrease) with the number of '1' in the binary
      representation of the given constant.

    :param n: The size of the quantum register we will add `number` to.
    :param number: The number to add to the given quantum register.
    :return: a routine adding number to the given quantum register.
    """
    add_routine = QRoutine()
    qubits = add_routine.new_wires(n)
    # Swap qubits endianness because qat's adder is in LSB
    add_routine.apply(classarith.add_const(n, number), qubits[::-1])
    return add_routine


@build_gate("sub_const", [int, int], arity=lambda n, num: n)
def sub_const_carry(n: int, number: int) -> QRoutine:
    r"""Subtract a constant to the given `n`\ -qubit register.

    The following properties are desirable ([ ] means that the property is not
    guaranteed to be fulfilled, [x] guarantees that the property is fulfilled
    by the current implementation):

    - [x] Looping underflow.
    - [ ] Exact results (no non-valid state with very small amplitudes
      appearing).
    - [x] Gate number increase (or decrease) with the number of '1' in the binary
      representation of the given constant.

    :param n: The size of the quantum register we will subtract `number` from.
    :param number: The number to subtract from the given quantum register.
    :return: a routine subtracting number from the given quantum register.
    """
    routine = QRoutine()
    sub_const_gate = add_const_carry(n, number).dag()
    qubits = routine.new_wires(sub_const_gate.arity)
    routine.apply(sub_const_gate, qubits)
    return routine


@build_gate("high_bit_compute", [int, int], arity=lambda n, rhs: 2 * n)
def high_bit_compute(n: int, rhs: int) -> QRoutine:
    r"""Compute the high-bit of the addition of the given quantum register and rhs.

    The implementation follows `<https://arxiv.org/pdf/1611.07995.pdf>`_ and the
    circuit is depicted in Figures 2 and 3.

    The returned routine needs :math:`2n` qubits organised as follow: ::

        |       |lhs>       |  |res>  |      |dirty>      |
        |   .   .   .   .   |    .    |   .   .   .   .   |
        |   0          n-1  |    n    |  n+1         2*n  |

    The :math:`\ket{\text{dirty}}` quantum register is a little special in this case
    as it does not necessarily need to be initialised to :math:`\ket{0}`. The
    :math:`\ket{\text{dirty}}` quantum register can be given in an arbitrary
    state and will be returned in the exact same state.

    :param n: The size of the quantum register storing the left-hand side of the
        addition.
    :param rhs: The number to add to the given quantum register.
    :return: a routine computing the high-bit of the operation lhs + rhs.
    """
    # The algorithm has been originally devised for little endian, the constant should
    # also be in little endian so we reverse its bits.
    rhsb = [b == "1" for b in bin(rhs)[2:].zfill(n)][::-1]
    rout = QRoutine()

    # The original circuit have been devised for a little-endian register.
    # In order to respect the big-endian convention of the qaths library and to keep
    # the algorithm as similar as possible to the one described in the paper, we
    # reverse the endianness on the indices right now.
    # The quantum register is, **WITHIN THIS PROCEDURE**, interpreted as a little-endian
    # quantum register.
    a = rout.new_wires(n)[::-1]
    res = rout.new_wires(1)
    g = rout.new_wires(n - 1)

    if n == 1:
        if rhsb[0]:
            rout.apply(CNOT, a[0], res)
        return rout

    # True routine:
    rout.apply(CNOT, g[-1], res)
    with rout.compute():
        for target_g_index in reversed(range(1, len(g))):
            a_index = target_g_index + 1
        if rhsb[a_index]:
            rout.apply(CNOT, a[a_index], g[target_g_index])
            rout.apply(X, a[a_index])
        rout.apply(CCNOT, g[target_g_index - 1], a[a_index], g[target_g_index])
        # For target_g_index == 1:
        if rhsb[1]:
            rout.apply(CNOT, a[1], g[0])
            rout.apply(X, a[1])
        # For target_g_index == 0:
        if rhsb[0]:
            rout.apply(CCNOT, a[0], a[1], g[0])
        # CNOT cascade
        for target_g_index in range(1, len(g)):
            rout.apply(
                CCNOT, g[target_g_index - 1], a[target_g_index + 1], g[target_g_index]
            )
    rout.apply(CNOT, g[-1], res)
    # Uncompute everything except res
    rout.uncompute()

    return rout
